import gensim
from gensim.models import Word2Vec
from gensim.models import KeyedVectors
import numpy as np
import tensorflow as tf
from tensorflow.keras import layers, models
import word2vector
import code2word as c2w
import os


def identify_file_type(file_path):
    _, ext = os.path.splitext(file_path)
    return ext.lower()

def predict_code_snippet(model, code_snippet, word2vec_model, max_length, vector_size):
    # 向量化代码片段
    snippet_matrix = np.zeros((max_length, vector_size))
    for i, token in enumerate(code_snippet):
        if i < max_length:
            snippet_matrix[i] = word2vec_model.wv[token] if token in word2vec_model.wv else np.zeros(vector_size)
    # 扩充矩阵使它符合模型输入预期
    snippet_matrix = snippet_matrix.reshape((1, max_length, vector_size, 1))
    # 使用模型进行预测
    prediction = model.predict(snippet_matrix)
    return prediction

def final_predict(file_path):

    file_type = identify_file_type(file_path)
    # 读取并转换代码段
    if file_type == ".py":
        code = c2w.py_code2word(file_path)
    elif file_type == ".java":
        code = c2w.go_Java_code2word(file_path)
    elif file_type == ".go":
        code = c2w.go_Java_code2word(file_path)
    else:
        print("Unsupported file type\n")
        exit()

    # 载入向量转换模型和检测模型
    model = tf.keras.models.load_model('./model/trojan_detector_cnn_word2vec.h5')
    word2vec_model = KeyedVectors.load('./model/word2vec.model')

    # 使用模型进行预测
    max_length = 800
    vector_size = 100

    prediction = predict_code_snippet(model, code, word2vec_model, max_length, vector_size)
    # 处理预测结果
    class_labels = ['normal', 'backdoor', 'dos_attack', 'ftp_exfiltration', 'keylogger', 'persistence_backdoor',
                    'ransomware', 'reverse_shell']
    predicted_probabilities = prediction[0]  # 获取预测的概率分布
    predicted_class_index = np.argmax(predicted_probabilities)  # 找到概率最高的类别索引
    predicted_class = class_labels[predicted_class_index]  # 对应的类别名称

    os.system('cls')
    # 打印每个类别及其概率
    print('''
                                             █                                                     ███                 
 ████▒           █                    █                                █████                         █      █          
 █  ▒█░          █                    █                                █   ▓█                        █      █          
 █   ▒█  ███   █████   ███    ▓██▒  █████  ███     ███   █▒██▒         █    █  ███   ▒███▒  █   █    █    █████  ▒███▒ 
 █    █ ▓▓ ▒█    █    ▓▓ ▒█  ▓█  ▓    █      █    █▓ ▓█  █▓ ▒█         █   ▒█ ▓▓ ▒█  █▒ ░█  █   █    █      █    █▒ ░█ 
 █    █ █   █    █    █   █  █░       █      █    █   █  █   █         █████  █   █  █▒░    █   █    █      █    █▒░   
 █    █ █████    █    █████  █        █      █    █   █  █   █         █  ░█▒ █████  ░███▒  █   █    █      █    ░███▒ 
 █   ▒█ █        █    █      █░       █      █    █   █  █   █         █   ░█ █         ▒█  █   █    █      █       ▒█ 
 █  ▒█░ ▓▓  █    █░   ▓▓  █  ▓█  ▓    █░     █    █▓ ▓█  █   █         █    █ ▓▓  █  █░ ▒█  █▒ ▓█    █░     █░   █░ ▒█ 
 ████▒   ███▒    ▒██   ███▒   ▓██▒    ▒██  █████   ███   █   █         █    ▒  ███▒  ▒███▒  ▒██▒█    ▒██    ▒██  ▒███▒
''')
    print("模型检测结果：")
    print("Prediction probabilities for each class:")
    for label, prob in zip(class_labels, predicted_probabilities):
        print(f"{label}: {prob:.6f}")

    print("Predicted class:", predicted_class)
    print(f"Probability of being a {predicted_class}: {predicted_probabilities[predicted_class_index]:.6f}")

    return class_labels, predicted_probabilities

def final_predict_files(file_path):

    file_type = identify_file_type(file_path)
    # 读取并转换代码段
    if file_type == ".py":
        code = c2w.py_code2word(file_path)
    elif file_type == ".java":
        code = c2w.go_Java_code2word(file_path)
    elif file_type == ".go":
        code = c2w.go_Java_code2word(file_path)
    else:
        print("Unsupported file type\n")
        exit()

    # 载入向量转换模型和检测模型
    model = tf.keras.models.load_model('./model/trojan_detector_cnn_word2vec.h5')
    word2vec_model = KeyedVectors.load('./model/word2vec.model')

    # 使用模型进行预测
    max_length = 800
    vector_size = 100

    prediction = predict_code_snippet(model, code, word2vec_model, max_length, vector_size)
    # 处理预测结果
    class_labels = ['normal', 'backdoor', 'dos_attack', 'ftp_exfiltration', 'keylogger', 'persistence_backdoor',
                    'ransomware', 'reverse_shell']
    predicted_probabilities = prediction[0]  # 获取预测的概率分布
    predicted_class_index = np.argmax(predicted_probabilities)  # 找到概率最高的类别索引
    predicted_class = class_labels[predicted_class_index]  # 对应的类别名称

    return class_labels, predicted_probabilities

if __name__ == "__main__":
    # 这里需要给出需要检测的代码处理后的形式
    file_path = 'hello.java'    # 给出实际待检测的文件路径
    final_predict(file_path)

    '''
    file_type = identify_file_type(file_path)
    # 读取并转换代码段
    if file_type == ".py":
        code = c2w.py_code2word(file_path)
    elif file_type == ".java":
        code = c2w.go_Java_code2word(file_path)
    elif file_type == ".go":
        code = c2w.go_Java_code2word(file_path)
    else:
        print("Unsupported file type\n")
        exit()

    # 载入向量转换模型和检测模型
    model = tf.keras.models.load_model('./model/trojan_detector_cnn_word2vec.h5')
    word2vec_model = KeyedVectors.load('./model/word2vec.model')

    # 使用模型进行预测
    max_length = 1000
    vector_size = 100

    prediction = predict_code_snippet(model, code, word2vec_model, max_length, vector_size)
    # 处理预测结果
    class_labels = ['normal', 'backdoor', 'dos_attack', 'ftp_exfiltration', 'keylogger', 'persistence_backdoor', 'ransomware', 'reverse_shell']
    predicted_probabilities = prediction[0]  # 获取预测的概率分布
    predicted_class_index = np.argmax(predicted_probabilities)  # 找到概率最高的类别索引
    predicted_class = class_labels[predicted_class_index]  # 对应的类别名称

    # 打印每个类别及其概率
    print("Prediction probabilities for each class:")
    for label, prob in zip(class_labels, predicted_probabilities):
        print(f"{label}: {prob:.6f}")

    print("Predicted class:", predicted_class)
    print(f"Probability of being a {predicted_class}: {predicted_probabilities[predicted_class_index]:.6f}")
    '''